import collections
import copy
import datetime
import pandas as pd
import pytz
import os
import json
from absl import app, flags
from typing import Dict, DefaultDict, List, Optional

import coin.base.database.util as db_util
from coin.base.datetime_util import to_timestamp_int
from coin.support.proto_log.logic.util import (
    enumerate_og_log_into_oginfos, read_og_log, OgInfo)
from coin.support.proto_log.app.og_log_replayer import OnLog
from coin.exchange.base.order_gateway_logger import parse_file_name_from_coin2
from coin.support.feed_tool.feed_stats.logic.util import fetch_git_commit_sha_datetime
from coin.support.pta.logic.pta_stats_calculator import (
    calculate_turnover, get_accounting_currency, calculate_turnover_in_usd)
from coin.proto.coin_strategy_summary_pb2 import (
    StrategySummary, StrategyStat, AccountStrategyStat)
from coin.proto.coin_request_pb2 import (
    ProcessInfoProto, CalculatorInfoProto, StrategyInfoProto)
from coin.tool.strat_monitor.util.monitor_util import load_pnl_adjust_config
from xunkemgmt_client.client.util.query_util import (
    query_accounts,
    query_applications,
    query_strategies,
    query_exchange_rates)
from coin2.service.account.account_info_service_pb2 import (
    QueryBalanceIntervalHistoriesRequestProto,
    QueryBalanceSnapshotsRequestProto,
)
from xunkemgmt_client.client.api_client import XunkemgmtClient
from xunkemgmt_client.support.pta.database.mysql_importer import PtaStatsImporter
from xunkemgmt_client.tool.slack_noti import send_to_slack

FLAGS = flags.FLAGS

flags.DEFINE_boolean('allow_empty_price', False, '')
flags.DEFINE_boolean('show_detail', False, 'show pnl report for per exchange')


def format_pnl_report(x: pd.DataFrame) -> str:
  if pd.isna(x):
    return ''
  elif isinstance(x, float):
    if abs(x) > 0.01:
      return "{:,.2f}".format(x)
    else:
      return "{:,.5f}".format(x)
  else:
    return x


def load_reserve_config(reserve_file: str) -> Dict:
  assert reserve_file and os.path.isfile(reserve_file), \
    'Reserve file does not exist: %s' % reserve_file
  with open(reserve_file, 'r') as open_file:
    reserve_info = json.load(open_file)
  reserve_info = collections.OrderedDict(sorted(reserve_info.items(), key=lambda x: x[0]))
  return reserve_info


def print_pnl_report(pnl_report_summary: DefaultDict[str, pd.DataFrame]) -> None:
  for strategy_name, pnl_report in copy.deepcopy(pnl_report_summary).items():
    pnl_report.dropna(how='all', axis=1, inplace=True)
    for col in pnl_report:
      if col[-1] in ('balance_est', 'reserve_est', 'diff_est', 'pnl'):
        pnl_report.loc['TOTAL', col] = pnl_report[col].sum(min_count=1)
    for col in pnl_report:
      pnl_report[col] = pnl_report[col].apply(format_pnl_report)
    print("Strategy: %s\n" % strategy_name)
    print(pnl_report.to_string())
    print("\n\n")


def gen_pnl_summary_print_str(pnl_report_summary: DefaultDict[str, pd.DataFrame]) -> str:
  result = ""
  pnl_summary = collections.defaultdict(pd.DataFrame)
  for strategy_name, pnl_report in copy.deepcopy(pnl_report_summary).items():
    for col in pnl_report:
      if col[-1] =='pnl':
        pnl_summary[col[0]].loc[strategy_name, 'pnl_in_usd'] = pnl_report[col].sum(min_count=1)
  for trading_date, pnl_report in pnl_summary.items():
    pnl_report.loc['TOTAL'] = pnl_report.sum(min_count=1)
    if pnl_report.dropna().empty:
      continue
    for col in pnl_report:
      pnl_report[col] = pnl_report[col].apply(format_pnl_report)
    result += "Trading date %s:\n" % trading_date
    result += pnl_report.to_string()
    result += "\n\n"
  return result


class CoinFOFPtaDumper(object):
  def __init__(self):
    self._reserve_info = load_reserve_config(
        '../../coin_deploy/support_pta/coinfof/reserve.json')
    self._pnl_adj = load_pnl_adjust_config(
        '../../coin_deploy/support_pta/coinfof/pnl_adj.json')
    self._account_df = query_accounts()
    self._application_df = query_applications()
    self._external_strategy_df = query_strategies(business_units=['CoinFOF'], active=True)
    self.pnl_report = collections.defaultdict(pd.DataFrame)
    self.turnover_report = collections.defaultdict(pd.DataFrame)

  def _print_per_exchange(self,
                          strategy_name: str,
                          account_ids: List[int],
                          trading_date: datetime.datetime,
                          coin_price: pd.Series
  ) -> None:
    account_per_strategy = self._account_df[self._account_df['account_id'].isin(account_ids)]
    for exchange in account_per_strategy['exchange'].drop_duplicates().sort_values():
      ids = account_per_strategy['account_id'][account_per_strategy['exchange'] == exchange].tolist()
      balance = self._query_balance_info(ids, trading_date)
      print(str(trading_date) + ' ' + strategy_name + ' ' + exchange + ':\n')
      res = pd.DataFrame()
      res['balance'] = balance.sum()
      res['price'] = coin_price
      res['value'] = res['balance'] * res['price']
      res.loc['TOTAL', 'value'] = res['value'].sum()
      for col in res.columns:
        res[col] = res[col].apply(format_pnl_report)
      print(res.to_string())
      print('\n\n')

  def _get_reserve_detail(self, trading_date: datetime.datetime, strategy_name: str) -> pd.Series:
    available_reserve_list = [
        date for date in self._reserve_info.keys() if date <= trading_date.strftime("%Y%m%d")
    ]
    reserve_info = None
    if available_reserve_list:
      for date in available_reserve_list:
        reserve_info = self._reserve_info[date].get(strategy_name) or reserve_info
    if reserve_info is None:
      return pd.Series(dtype='float64')
    reserve = pd.Series(
        [reserve['reserve'] for reserve in reserve_info],
        index=[reserve['accounting_currency'] for reserve in reserve_info])
    return reserve

  def _get_pnl_adj_detail(self, trading_date: datetime.datetime, strategy_name: str) -> pd.Series:
    available_pnl_adj = self._pnl_adj.get(trading_date.strftime("%Y%m%d"))
    if available_pnl_adj:
      available_pnl_adj = pd.DataFrame(available_pnl_adj)
      available_pnl_adj = available_pnl_adj[available_pnl_adj['strategy_name'] == strategy_name]
    else:
      available_pnl_adj = pd.DataFrame(columns=['accounting_currency', 'pnl_adj', 'strategy_name'])
    available_pnl_adj = pd.Series(
        available_pnl_adj['pnl_adj'].tolist(),
        index=available_pnl_adj['accounting_currency'], dtype='float64')
    return available_pnl_adj

  def _query_balance_info(self, account_ids: List[int], trading_date: Optional[datetime.datetime]=None
  ) -> pd.DataFrame:
    with XunkemgmtClient() as client:
      if trading_date is not None:
        request = QueryBalanceIntervalHistoriesRequestProto(
            account_ids = account_ids,
            start_timestamp = to_timestamp_int(trading_date),
            end_timestamp = to_timestamp_int(trading_date + datetime.timedelta(days=1)),
        )
        response_list = client.query_balance_interval_histories(request).histories
      else:
        request = QueryBalanceSnapshotsRequestProto(
            account_ids = account_ids,
        )
        response_list = client.query_balance_snapshots(request).snapshots
    balance_list = []
    for resp in response_list:
      account_id = resp.balance_info.account_id
      for balance in resp.each_balance:
        balance_list.append([account_id, balance.currency, balance.total])
    balance_df = pd.DataFrame(balance_list, columns=['account_id', 'currency', 'total'])
    balance_df = balance_df.pivot(index='account_id', columns='currency', values='total')
    return balance_df

  def _gen_estimate_report_impl(self,
                                trading_date: datetime.datetime,
                                balance: pd.Series,
                                reserve_detail: pd.Series,
                                pnl_adj_detail: pd.Series,
                                coin_price: pd.Series
  ) -> pd.DataFrame:
    df = pd.DataFrame()
    df['balance'] = balance
    df['reserve'] = reserve_detail
    df['pnl_adj'] = pnl_adj_detail
    df = df.join(coin_price.to_frame(name='price'), how='inner')
    if not FLAGS.allow_empty_price:
      assert not df['price'].hasnans
    df['balance_est'] = df['balance'] * df['price']
    df['reserve_est'] = df['reserve'] * df['price']
    df['pnl_adj_est'] = df['pnl_adj'] * df['price']
    df['diff_est'] = df['balance_est'].subtract(df['reserve_est'], fill_value=0)
    df.columns = pd.MultiIndex.from_product(
        [[str(trading_date)], df.columns], names=['Trading_date', ''])
    return df

  def _cal_daily_pnl(self,
                     pnl_report: DefaultDict[str, pd.DataFrame],
                     trading_date: datetime.datetime,
                     strategy_name: str
  ) -> DefaultDict[str, pd.DataFrame]:
    strat_pnl_report = pnl_report[strategy_name]
    try:
      this_acc_pnl = strat_pnl_report[(str(trading_date), 'diff_est')]
      last_acc_pnl = strat_pnl_report[(str(trading_date - datetime.timedelta(days=1)), 'diff_est')]
      this_pnl_adj = strat_pnl_report[(str(trading_date), 'pnl_adj_est')]
      if this_acc_pnl.dropna().empty or last_acc_pnl.dropna().empty:
        pnl = pd.Series()
      else:
        pnl = this_acc_pnl.fillna(0) - last_acc_pnl.fillna(0) + this_pnl_adj.fillna(0)
    except KeyError:
      pnl = pd.Series(dtype='float64')
    except Exception as e:
      pnl = pd.Series(dtype='float64')
      print("Fail to cal pnl in %s for %s :" % (trading_date, strategy_name), type(e), e)
    finally:
      strat_pnl_report[(str(trading_date), 'pnl')] = pnl
    return pnl_report

  def _gen_daily_balance_report(self, trading_date: datetime.datetime
  ) -> Dict[str, pd.DataFrame]:
    balance_report = {}
    external_account_ids = self._application_df['account_id'][
        self._application_df['strategy_id'].isin(
        self._external_strategy_df['strategy_id'])].tolist()
    total_balance = self._query_balance_info(external_account_ids, trading_date)
    for strategy_name in self._external_strategy_df['strategy_name']:
      account_ids = self._application_df['account_id'][
        self._application_df['strategy_name'] == strategy_name].tolist()
      balance = total_balance.loc[[id for id in total_balance.index if id in account_ids
                                  ]].T.dropna(how='all').sum(axis=1)
      reserve_detail = self._get_reserve_detail(trading_date, strategy_name)
      pnl_adj_detail = self._get_pnl_adj_detail(trading_date, strategy_name)
      coin_list = list(set(balance.index.tolist() + list(reserve_detail.index)))
      coin_price = query_exchange_rates(coin_list, 'USD', trading_date=trading_date)
      coin_price = pd.Series(coin_price, index=coin_list)
      if not FLAGS.allow_empty_price:
        if coin_price.hasnans:
          raise ValueError("No price for %s in %s" % (
              coin_price[coin_price.isna()].index.tolist(), trading_date))
      estimate_df = self._gen_estimate_report_impl(
          trading_date, balance, reserve_detail, pnl_adj_detail, coin_price)
      balance_report[strategy_name] = estimate_df
      if FLAGS.show_detail:
        self._print_per_exchange(strategy_name, account_ids, trading_date, coin_price)
    return balance_report

  def _gen_daily_turnover_report(self, trading_date: datetime.datetime
  ) -> Dict[str, pd.DataFrame]:
    turnover_report = {}
    for strategy_name in self._external_strategy_df['strategy_name']:
      account_df = self._application_df[
        self._application_df['strategy_name'] == strategy_name]
      acct_info_from_archive = enumerate_og_log_into_oginfos(
        FLAGS.account_log_dir, trading_date.strftime('%Y%m%d'), parse_file_name_from_coin2)
      acct_info_df = pd.DataFrame(acct_info_from_archive)
      assert not acct_info_df.empty, f'missing account proto log {trading_date} {strategy_name}'
      account_df = account_df.rename(columns={
        'market_type_acct': 'market_type',
        'exchange_acct': 'exchange',
        'owner_acct': 'owner'})
      account_df = account_df.merge(acct_info_df, on=['market_type', 'exchange', 'owner'])
      order_fill_df_list = []
      for _, row in account_df.iterrows():
        order_fill_df = dump_order_fill(
          trading_date, row['market_type'], row['exchange'], row['owner'], row['machine'])
        order_fill_df['account_id'] = row['account_id']
        order_fill_df_list.append(order_fill_df)
      order_fill_df_list = [order_fill_df for order_fill_df in order_fill_df_list if not order_fill_df.empty]
      if not order_fill_df_list:
        continue
      order_fill_df = pd.concat(order_fill_df_list)
      order_fill_df = order_fill_df.rename(columns={'fill_qty': 'volume_mark'})
      order_fill_df['accounting_currency'] = order_fill_df.apply(
        lambda x: get_accounting_currency(x['market_type'], x['exchange_type'], x['symbol'], trading_date), axis=1)
      order_fill_df['price_in_usd'] = query_exchange_rates(
        order_fill_df['accounting_currency'].tolist(), 'USD', trading_date)
      for fill_type in ('maker', 'taker'):
        order_fill_df[f'volume_{fill_type}_mark'] = order_fill_df.apply(
          lambda x: x['volume_mark'] if x['fill_type'] == f'{fill_type.upper()}_FILL_TYPE' else 0, axis=1)
      for fill_type in ('mark', 'maker_mark', 'taker_mark'):
        order_fill_df[f'turnover_{fill_type}'] = order_fill_df.apply(
          lambda x: calculate_turnover(
              market_type=x['market_type'],
              exchange=x['exchange_type'],
              symbol=x['symbol'],
              volume=x[f'volume_{fill_type}'],
              price=x['fill_price'],
              current_datetime=trading_date), axis=1)
        order_fill_df[f'turnover_{fill_type}_in_usd'] = order_fill_df.apply(
          lambda x: calculate_turnover_in_usd(
              market_type=x['market_type'],
              exchange=x['exchange_type'],
              symbol=x['symbol'],
              volume=x[f'volume_{fill_type}'],
              turnover=x[f'turnover_{fill_type}'],
              turnover_price_in_usd=x['price_in_usd'],
              current_datetime=trading_date), axis=1)
      order_fill_df = order_fill_df.reset_index()
      order_fill_df.columns = pd.MultiIndex.from_product(
          [[str(trading_date)], order_fill_df.columns], names=['Trading_date', ''])
      turnover_report[strategy_name] = order_fill_df
    return turnover_report

  def get_pnl_report(self, start_date: datetime.datetime, end_date: datetime.datetime
  ) -> DefaultDict[str, pd.DataFrame]:
    trading_dates = [dt.date() for dt in pd.date_range(start_date - datetime.timedelta(days=1), end_date)]
    pnl_report = collections.defaultdict(pd.DataFrame)
    for trading_date in trading_dates:
      balance_report = self._gen_daily_balance_report(trading_date)
      for strategy_name in balance_report:
        pnl_report[strategy_name] = pd.concat(
          [pnl_report[strategy_name], balance_report[strategy_name]], axis=1)
        pnl_report = self._cal_daily_pnl(pnl_report, trading_date, strategy_name)
    self.pnl_report = pnl_report
    return pnl_report

  def get_turnover_report(self, start_date: datetime.datetime, end_date: datetime.datetime
  ) -> DefaultDict[str, pd.DataFrame]:
    trading_dates = [dt.date() for dt in pd.date_range(start_date, end_date)]
    turnover_report = collections.defaultdict(pd.DataFrame)
    for trading_date in trading_dates:
      daily_turnover_report = self._gen_daily_turnover_report(trading_date)
      for strategy_name in daily_turnover_report:
        turnover_report[strategy_name] = pd.concat(
          [turnover_report[strategy_name], daily_turnover_report[strategy_name]], axis=1)
    return turnover_report

  def _gen_calculator_proto(self) -> CalculatorInfoProto:
    git_commit_datetime, git_commit_sha = fetch_git_commit_sha_datetime('HEAD')
    git_commit_datetime = int(git_commit_datetime.replace(tzinfo=pytz.UTC).timestamp() * 10 ** 9)
    calculator_proto = CalculatorInfoProto(calculator_name=os.path.basename(__file__),
                                           git_commit_sha=git_commit_sha,
                                           git_commit_timestamp=git_commit_datetime)
    return calculator_proto

  def get_strat_summaries(self, start_date: datetime.datetime, end_date: datetime.datetime
  ) -> List[StrategySummary]:
    pnl_report_summary = self.pnl_report if self.pnl_report else \
      self.get_pnl_report(start_date, end_date)
    turnover_report_summary = self.turnover_report if self.turnover_report else \
      self.get_turnover_report(start_date, end_date)
    summary_list = []
    calculator_proto = self._gen_calculator_proto()
    for strategy_name, pnl_report in pnl_report_summary.items():
      for trading_date in pnl_report.columns.levels[0]:
        daily_pnl = pnl_report[(str(trading_date), 'pnl')]
        total_pnl = daily_pnl.sum(min_count=1)
        if pd.isna(total_pnl):
          continue
        strat_stat = StrategyStat(pnl_mark=total_pnl,
                                  pnl_mark_in_usd=total_pnl,
                                  pnl_mark_from_balance=total_pnl,
                                  pnl_mark_from_balance_in_usd=total_pnl,
                                  accounting_currency='USD')
        stat_per_account = collections.defaultdict(dict)
        if strategy_name in turnover_report_summary:
          daily_turnover = turnover_report_summary[strategy_name].get(str(trading_date), pd.DataFrame())
          if not daily_turnover.empty:
            daily_turnover = daily_turnover[
              ['account_id', 'symbol', 'volume_maker_mark', 'volume_taker_mark', 'volume_mark',
              'turnover_maker_mark_in_usd', 'turnover_taker_mark_in_usd', 'turnover_mark_in_usd']]
            daily_turnover = daily_turnover.groupby(['account_id', 'symbol']).sum()
            daily_turnover['accounting_currency'] = 'USD'
            turnover_dict = daily_turnover.to_dict('index')
            for key, stats in turnover_dict.items():
              stat = stat_per_account[int(key[0])]
              stat.update({key[1]: StrategyStat(**stats)})
            stat_per_account = {key: AccountStrategyStat(stat_per_symbol=stats_per_symbol)
                                for key, stats_per_symbol in stat_per_account.items()}
            strat_stat.turnover_maker_mark_in_usd = daily_turnover['turnover_maker_mark_in_usd'].sum()
            strat_stat.turnover_taker_mark_in_usd = daily_turnover['turnover_taker_mark_in_usd'].sum()
            strat_stat.turnover_mark_in_usd = daily_turnover['turnover_mark_in_usd'].sum()

        strategy = self._external_strategy_df[
          self._external_strategy_df['strategy_name'] == strategy_name].iloc[0]
        strategy_info = StrategyInfoProto(
            strategy_id=strategy['strategy_id'],
            strategy_name=strategy['strategy_name'],
            strategy_group=strategy['strategy_group']
        )
        trading_date = datetime.datetime.strptime(trading_date, "%Y-%m-%d").date()
        start_dt = datetime.datetime.combine(
          trading_date, datetime.time.min).replace(tzinfo=pytz.UTC)
        end_dt = datetime.datetime.combine(
          trading_date, datetime.time.max).replace(tzinfo=pytz.UTC)
        process_info = ProcessInfoProto(
            trading_date=int(start_dt.date().strftime('%Y%m%d')),
            start_timestamp=int(start_dt.timestamp() * 1e+9),
            end_timestamp=int(end_dt.timestamp() * 1e+9),
            machine='',
        )
        strat_summary_proto = StrategySummary(
            strategy=strategy_info,
            calculator=calculator_proto,
            process=process_info,
            stat_per_quote={'USD': strat_stat},
            stat_per_account=stat_per_account,
        )
        summary_list.append(strat_summary_proto)
    return summary_list


def dump_order_fill(trading_date: datetime.datetime,
                    market_type: str,
                    exchange: str,
                    owner: str,
                    machine: str
) -> pd.DataFrame:
  td_str = trading_date.strftime('%Y%m%d')
  filename = f"/tmp/{market_type}_{exchange}_{owner}_{td_str}.json"
  start_dt = datetime.datetime.combine(
    trading_date, datetime.time.min).replace(tzinfo=pytz.UTC)
  end_dt = datetime.datetime.combine(
    trading_date, datetime.time.max).replace(tzinfo=pytz.UTC)
  onlog = OnLog(filename)
  og_info = OgInfo(market_type=market_type, exchange=exchange, owner=owner,
                  machine=machine, trading_date=None)
  read_og_log(og_info=og_info, start_time=start_dt, end_time=end_dt,
              callback=onlog.on_log, root_dir=FLAGS.account_log_dir)
  onlog.close()
  order_fill_dict = {}
  for pb in open(filename, 'r'):
    account_info = json.loads(pb)
    if account_info['type'] == 'ORDER_EVENT':
      order_event = account_info['event']
      order_fill_dict[
        (order_event['external_order_id'], order_event['proc_order_id'], order_event['symbol'])
      ] = account_info['event']
  order_fill_df = pd.DataFrame(order_fill_dict.values())
  return order_fill_df


def main(_):
  pd.set_option('display.width', 200)
  pd.set_option('max_colwidth', 200)
  max_date = datetime.datetime.utcnow().date()
  if FLAGS.start_date is None and FLAGS.end_date is None:
    start_date = max_date - datetime.timedelta(days=1)
    end_date = start_date
  else:
    start_date = datetime.datetime.strptime(FLAGS.start_date, '%Y%m%d').date()
    end_date = datetime.datetime.strptime(FLAGS.end_date, '%Y%m%d').date()
  assert start_date <= end_date <= max_date, (start_date, end_date, max_date)
  pta_dumper = CoinFOFPtaDumper()
  pnl_report_summary = pta_dumper.get_pnl_report(start_date, end_date)
  print_pnl_report(pnl_report_summary)
  pnl_summary_str = gen_pnl_summary_print_str(pnl_report_summary)
  print(pnl_summary_str)
  if FLAGS.slack_receiver:
    msg = "```" + pnl_summary_str + "```"
    send_to_slack(msg, FLAGS.slack_receiver, 'msg')
  pta_db_config = db_util.read_db_config(FLAGS.pta_db_config) if FLAGS.pta_db_config else None
  if pta_db_config:
    importer = PtaStatsImporter(db_config=pta_db_config)
    strat_summary_proto_list = pta_dumper.get_strat_summaries(start_date, end_date)
    for strat_summary_proto in strat_summary_proto_list:
      strat_summary_proto.source_info = pta_db_config['database']
      importer.insert_strategy_pta_stats(strat_summary_proto)


if __name__ == '__main__':
  flags.DEFINE_string('start_date', None, '%Y%m%d')
  flags.DEFINE_string('end_date', None, '%Y%m%d')
  flags.DEFINE_string('slack_receiver', None, 'slack channel')
  flags.DEFINE_string('pta_db_config', None, 'path of pta db config')
  flags.DEFINE_string('account_log_dir',
                      '/remote/iosg/strat-1/buckets/log.raw.coin/live/account_proto_log',
                      'account proto log dir')
  app.run(main)
